use std::{env, fmt, fs};
use std::collections::HashMap;
use serde::de::{MapAccess, Visitor};
use serde::Deserialize;

#[derive(Debug)]
struct MergedMap(HashMap<String, Vec<String>>);

impl<'a> Deserialize<'a> for MergedMap {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
        where
            D: serde::Deserializer<'a> {
                deserializer.deserialize_map(MergedMapVisitor)
    }
}

struct MergedMapVisitor;
impl<'de> Visitor<'de> for MergedMapVisitor {
    type Value = MergedMap;
    fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
        formatter.write_str("a map with potentially duplicate keys")
    }
    fn visit_map<A>(self, mut map: A) -> Result<Self::Value, A::Error>
    where
        A: MapAccess<'de>,
    {
        let mut merged_map: HashMap<String, Vec<String>> = HashMap::new();

        while let Some((key, value)) = map.next_entry::<String, Vec<String>>()? {
            merged_map
                .entry(key)
                .or_insert_with(Vec::new) // 如果键不存在，则插入一个空的 Vec
                .extend(value); // 合并值
        }
        Ok(MergedMap(merged_map))
    }
}

fn name_to_index(name: &str, citys: &mut HashMap<String, usize>) -> usize {
    if citys.contains_key(name) {
        *citys.get(name).unwrap()
    } else {
        let index = citys.len();
        citys.insert(name.to_string(), index);
        index
    }
}

fn find(parent: &mut Vec<usize>, x: usize) -> usize {
    if parent[x] != x {
        parent[x] = find(parent, parent[x]);
    }
    parent[x]
}


fn sub_count_provinces(group: &MergedMap) -> usize {
    let mut cities = HashMap::new();
    // 获取城市索引
    for v_n in &group.0 {
        name_to_index(&v_n.0, &mut cities);
        for u_n in v_n.1.iter() {
            name_to_index(u_n, &mut cities);
        }
    }
    // 初始化并查集
    let mut parent = vec![0; cities.len()];
    for i in 0..parent.len() {
        parent[i] = i;
    }
    // 合并城市
    for v_n in &group.0 {
        let v = *cities.get(v_n.0).unwrap();
        for u_n in v_n.1.iter() {
            let u = *cities.get(u_n).unwrap();
            if u == v {
                continue;
            }
            let root_v = find(&mut parent, v);
            let root_u = find(&mut parent, u);
            if root_v != root_u {
                parent[root_v] = root_u;
            }
        }
    }
    // 计算省份数量
    let mut count = 0;
    for i in 0..parent.len() {
        if parent[i] == i {
            count += 1;
        }
    }
    count
}

pub fn count_provinces() -> String {
    // 打印当前工作目录
    println!("Current directory is: {}", env::current_dir().unwrap().display());
    let contents = fs::read_to_string("district.json").unwrap();
    let data: HashMap<String, MergedMap> = serde_json::from_str(&contents).unwrap();
    let mut counts = vec![0; data.len()];
    // 获取城市索引
    for (index, group) in data.iter() {
        counts[index.parse::<usize>().unwrap()-1] = sub_count_provinces(group);
    }
    counts.iter().map(|x| x.to_string()).collect::<Vec<String>>().join(",")
}
